#!/usr/bin/env python

# -------
# Imports:
# -------
# Project configuration
from config import *

# Ros imports:
import roslib; roslib.load_manifest(NODE_NAME)
import rospy

# Other imports:
import os
from particle import Particle
from random import random
	
# ----------------
class AncestryTree:
# ----------------
	# ------------------------------------
	def __init__(self, numParticles, root):
	# ------------------------------------
		self.numParticles = numParticles			# The total number of particles
		self.particleLeafs = []						# The leaf particles (currently in consideration)
		self.root = root							# Our root particle
		self.highestProbLeaf = self.root			# What is currently the most likely particle

		# Populate the ancestry tree by running importance sampling 
		self.particleLeafs.append(self.root)
		self.importanceSampling()
		
	# --------------------------
	def importanceSampling(self):	
	# --------------------------
		# Make sure the particle probabilities are normalized
		self.normalizeProbabilities(self.particleLeafs)

		# Generate the cumulative probability distribution (list)
		cumulativeDistribution = []	
		cumulativeProbability = 0.0
		for p in self.particleLeafs:
			cumulativeProbability += p.probability
			cumulativeDistribution.append(cumulativeProbability)

		# Generate our NUM_PARTICLES new particles
		tempParticleLeafs = [] 
		for i in range(NUM_PARTICLES):
			# Generate a random number between 0 and 1
			randNum = random()

			# Search the cumultive distribution for the position in which our random number would fall
			particleIndex = self.numParticles - 1
			for (index, probability) in enumerate(cumulativeDistribution):
				if randNum <= probability:
					particleIndex = index
					break
			particleToSample = self.particleLeafs[particleIndex]

			# Sample the particle and add it to our new list of leafs
			tempParticleLeafs.append(particleToSample.sample())		

		# Normalize
		tempParticleLeafs = self.normalizeProbabilities(tempParticleLeafs)

		# Loop through all of the old particle leafs and prune the branches that don't
		# have any children (and therefore weren't sampled)
		for particle in self.particleLeafs:
			# Pruning is done recursively in the particles themselves...
			particle.prune()

		# Replace our leafs
		self.particleLeafs = tempParticleLeafs

		# Now collapse the tree (also done recursively)
		for particle in self.particleLeafs:
			particle.collapseUp()

		# Find the highest probability leaf
		self.findHighestProbLeaf()

	# ---------------------------
	def findHighestProbLeaf(self):
	# ---------------------------
		# Find the highest probability leaf
		self.highestProbLeaf = self.particleLeafs[0]
		for particle in self.particleLeafs:
			if particle.probability > self.highestProbLeaf.probability:
				self.highestProbLeaf = particle
		return self.highestProbLeaf
	
	# ----------------------------------------------------
	def normalizeProbabilities(self, particlesToNormalize):
	# ----------------------------------------------------
		# Sum the total probability of all the particles
		totalProbability = 0
		for particle in particlesToNormalize:
			totalProbability += particle.probability
		
		# Normalize probabilities
		for particle in particlesToNormalize:
			particle.probability /= totalProbability

		return particlesToNormalize

	# --------------------------------------
	def sortParticles(self, particlesToSort):
	# --------------------------------------
		probability = lambda particle: particle.probability
		particlesToSort.sort(key=probability, reverse = True)

	# ------------------------------------------
	def visualize(self, fileName, generateImage):
	# ------------------------------------------
		# Write the ancestry tree toString output to a temp file
		file = open(fileName + '.dot', 'w')
		file.write(self.visualizeToString())
		file.close()
	
		# Generate the image file using EOG	
		if generateImage:
			os.system('dot -Tpng -o %s.png %s.dot' % (fileName, fileName))
			os.system('eog %s.png' % fileName)

	# -------------------------
	def visualizeToString(self):
	# -------------------------
		# Create a set of particles we've seen
		visited = set()

		# Visualize the ancestry tree by generating a graphviz file
		toReturn = 'digraph ancestry_tree{\n'

		# Loop through all of the particles
		for particle in self.particleLeafs:
				# Handle the leafs separately
				visited.add(particle)	
				toReturn += 'particle%i[label="id: %i, probability:%f" color="red"]\n' % (particle.particleId, particle.particleId, particle.probability)
				if particle.parent is not None:
					toReturn += 'particle%i -> particle%i;\n' %  (particle.particleId, particle.parent.particleId)

				# Walk up the chain until we get to something we've seen, or we walk off the top
				ptr = particle.parent
				while ptr is not None and ptr.particleId not in visited:
					# We've now visited ptr
					visited.add(ptr.particleId)					
	
					# Establish the link to the parent
					toReturn += 'particle%i[label="id: %i, probability:%f"];' % (ptr.particleId, ptr.particleId, ptr.probability)
					if ptr.parent is not None:
						toReturn += 'particle%i -> particle%i;' %  (ptr.particleId, ptr.parent.particleId)
		
					# Move the pointer up
					ptr = ptr.parent 
		toReturn += '}'

		return toReturn;

	# -------------------------
	def getNumChildrenSet(self):
	# -------------------------
		s = {}
		for particle in self.particleLeafs:
			s[particle.particleId] = particle.numChildren
			s[particle.parent.particleId] = particle.parent.numChildren
		return s

	# ------------------------
	def __getitem__(self, key):
	# ------------------------
	# Allow access to the particle *leafs*
		return self.particleLeafs[key]

	# ---------------
	def __str__(self):
	# ---------------
		toReturn = 'Leafs:\n'
		for particle in self.particleLeafs:
			toReturn += str(particle) + '\n'
		return toReturn

